// This file is a part of Funky C++.
// Copyright (C) 2011  Valentin David
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.
#ifndef FFT_HH
# define FFT_HH

# include <signal/bus.hh>
# include <complex>
# include <fast/static_bin_log.hh>
# include <fast/invert_bits.hh>

namespace funky {

  template <typename T>
  struct out_type {
  };

  template <typename T, size_t N>
  struct out_type<bus<T,N> > {
    typedef T real;
    typedef std::complex<T> sample;
    typedef bus<std::complex<T>,N> type;
    static type convert(const bus<T,N>& in) {
      type ret;
      for (size_t i = 0; i < N; ++i) {
	ret[i] = std::complex<T>{in[i], 0.};
      }
      return std::move(ret);
    }
  };

  template <typename T, size_t N>
  struct out_type<bus<std::complex<T>, N> > {
    typedef T real;
    typedef std::complex<T> sample;
    typedef bus<std::complex<T>,N> type;
    static const type& convert(const type& in) {
      return in;
    }
    static type convert(const type&& in) {
      return std::move(in);
    }
  };

  template <typename Signal>
  struct dft_signal {
  private:
    typedef typename signal_traits<Signal>::sample bus_sample;
    static const size_t N = bus_sample::size;
    typedef typename out_type<bus_sample>::sample cplx_sample;
    typedef typename bus_sample::sample sample;
    typedef typename out_type<bus_sample>::real real_sample;

  public:
    static_assert((1 << fast::static_bin_log
		   <size_t, N>::value) == N,
		  "N is not a power of two");

    typedef typename out_type<bus_sample>::type out_sample;

  private:
    static constexpr size_t shift =
      8*sizeof(size_t) - fast::static_bin_log<size_t, N>::value;

    template <size_t Begin, size_t End,
	      typename = typename std::enable_if<Begin == End>::type,
	      size_t invert = (fast::static_invert_bits<size_t, Begin>::value
			       >>shift),
              typename = typename std::enable_if<(invert > Begin)>::type>
    static void reorder(out_sample& v) {
      std::swap(v[invert], v[Begin]);
    }

    template <size_t Begin, size_t End,
	      typename = typename std::enable_if<Begin == End>::type,
	      size_t invert = (fast::static_invert_bits<size_t, Begin>::value
			       >>shift),
	      typename = typename std::enable_if<!(invert > Begin)>::type,
              typename = void>
    static void reorder(out_sample&) {
    }

    template <size_t Begin, size_t End,
	      typename = typename std::enable_if<(Begin > End)>::type>
    static void reorder(out_sample&) {
    }

    template <size_t Begin, size_t End,
	      typename = typename std::enable_if<(Begin < End)>::type,
 	      typename = void>
    static void reorder(out_sample& v) {
      static constexpr size_t middle = Begin + (End - Begin)/2;
      reorder<Begin,  middle>(v);
      reorder<middle+1, End>(v);
    }


    static void dft_inplace(out_sample& ret, bool inverse) {
      reorder<0, N-1>(ret);

      static const cplx_sample I{0.,1.};
      static bus<cplx_sample, N> mult;
      static bool init = false;
      if (!init) {
	size_t n = 0;
	for (size_t step = 1; step < N; step <<= 1, ++n) {
	  real_sample delta = M_PI / step;
	  real_sample sine = sin(delta * .5f);
	  mult[n] = cplx_sample{-2.f * sine * sine, (real_sample)sin(delta)};
	}
	init = true;
      }

      size_t n = 0;
      for (size_t step = 1; step < N; step <<= 1, ++n) {
	size_t jump = step << 1;
	cplx_sample factor{1., 0.};
	for (size_t group = 0; group < step; ++group) {
	  for (size_t pair = group; pair < N; pair += jump) {
	    size_t match = pair + step;
	    cplx_sample p = factor*ret[match];
	    ret[match] = ret[pair] - p;
	    ret[pair] += p;
	  }
	  factor = (inverse?std::conj(mult[n]):mult[n]) * factor + factor;
	}
      }
    }

    static out_sample dft(bus_sample&& in, bool inverse) {
      out_sample ret{out_type<bus_sample>::convert(std::move(in))};
      dft_inplace(ret, inverse);
      return ret;
    }

    static out_sample dft(const bus_sample& in, bool inverse) {
      out_sample ret{out_type<bus_sample>::convert(in)};
      dft_inplace(ret, inverse);
      return ret;
    }

    Signal signal;
    bool inverse;

  public:
    explicit dft_signal(Signal& signal, bool inverse):
      signal(signal), inverse(inverse) {}

    explicit dft_signal(Signal&& signal, bool inverse):
      signal(std::move(signal)), inverse(inverse) {}

    dft_signal(dft_signal&& other):
      signal(std::move(other.signal)), inverse(other.inverse) {}

    dft_signal(const dft_signal&) = default;

    dft_signal& operator=(const dft_signal& other) {
      return *this = dft_signal(other);
    }

    dft_signal& operator=(dft_signal&& other) throw() {
      std::swap(signal, other.signal);
      std::swap(inverse, other.inverse);
      return *this;
    }

    dft_signal& operator++() {
      ++signal;
      return *this;
    }

    dft_signal operator++(int) {
      dft_signal ret{*this};
      ++*this;
      return ret;
    }

    out_sample operator*() const {
      return dft(*signal, inverse);
    }
  };

  template <typename T>
  struct signal_traits<dft_signal<T> > {
    typedef typename dft_signal<T>::out_sample sample;
    typedef size_t duration;
  };

  template <typename T, typename RT = typename std::decay<T>::type,
	    typename = typename signal_traits<RT>::sample>
  dft_signal<RT>
  dft(T&& t, bool inverse = false) {
    return dft_signal<RT>{std::forward<T>(t), inverse};
  }

}

#endif
